import javax.crypto.*;
import javax.crypto.spec.SecretKeySpec;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

public class AES {
    private static final String ALGORITHM = "AES";
    private static final int KEY_SIZE = 128; // 可选：128、192、256
    private static final String CIPHER_TRANSFORMATION = "AES/ECB/PKCS5Padding";

    /*    public static void main(String[] args) throws Exception {

    // 输入输出文件名
            //加密文件路径
            String inputFile = "/Users/ni/Desktop/勒索病毒测试/勒索病毒测试.pptx";
            //加密好的文件路径
            String encryptedFile = "/Users/ni/Desktop/勒索病毒测试/勒索病毒测试1.pptx";
            //解密文件路径
            String decryptedFile = "/Users/ni/Desktop/勒索病毒测试/勒索病毒测试2.pptx";

    // 生成AES密钥
    //        byte[] aesKeyBytes = generateAesKey();
    //        SecretKey aesKey = new SecretKeySpec(aesKeyBytes, ALGORITHM);
    //        System.out.println(aesKey.hashCode());

            //从文件中读取密钥
            File keyFileObj = new File("./key/key.dat");
            FileInputStream keyIn = new FileInputStream(keyFileObj);
            byte[] aesKeyBytes = new byte[(int) keyFileObj.length()];
            keyIn.read(aesKeyBytes);
            keyIn.close();
            SecretKey aesKey = new SecretKeySpec(aesKeyBytes, "AES");

    // 加密文件
            encryptFile(inputFile, encryptedFile, aesKey);

    // 解密文件
            decryptFile(encryptedFile, decryptedFile, aesKey);
        }

        private static byte[] generateAesKey() throws NoSuchAlgorithmException {
            KeyGenerator keyGenerator = KeyGenerator.getInstance(ALGORITHM);
            keyGenerator.init(KEY_SIZE);
            SecretKey secretKey = keyGenerator.generateKey();
            return secretKey.getEncoded();
        }

        private static void encryptFile(String inputFile, String outputFile, SecretKey key)
                throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException,
                IOException {
            Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORMATION);
            cipher.init(Cipher.ENCRYPT_MODE, key);
            FileInputStream inputStream = new FileInputStream(inputFile);
            FileOutputStream outputStream = new FileOutputStream(outputFile);
            CipherOutputStream cipherOutputStream = new CipherOutputStream(outputStream, cipher);
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                cipherOutputStream.write(buffer, 0, bytesRead);
            }
            cipherOutputStream.close();
            inputStream.close();
        }

        private static void decryptFile(String inputFile, String outputFile, SecretKey key)
                throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException,
                IOException {
            Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORMATION);
            cipher.init(Cipher.DECRYPT_MODE, key);
            FileInputStream inputStream = new FileInputStream(inputFile);
            CipherInputStream cipherInputStream = new CipherInputStream(inputStream, cipher);
            FileOutputStream outputStream = new FileOutputStream(outputFile);
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = cipherInputStream.read(buffer)) != -1) {
                outputStream.write(buffer, 0, bytesRead);
            }
            outputStream.close();
            cipherInputStream.close();
        }*/
    public static SecretKey getAesKey() throws IOException {
        File keyFileObj = new File("./key/key.dat");
        FileInputStream keyIn = new FileInputStream(keyFileObj);
        byte[] aesKeyBytes = new byte[(int) keyFileObj.length()];
        keyIn.read(aesKeyBytes);
        keyIn.close();
        SecretKey aesKey = new SecretKeySpec(aesKeyBytes, "AES");
        return aesKey;
    }

    //加密文件
    public static void encryptFile(String inputFile, String outputFile, SecretKey key)
            throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException,
            IOException {
        Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORMATION);
        cipher.init(Cipher.ENCRYPT_MODE, key);
        FileInputStream inputStream = new FileInputStream(inputFile);
        FileOutputStream outputStream = new FileOutputStream(outputFile);
        CipherOutputStream cipherOutputStream = new CipherOutputStream(outputStream, cipher);
        byte[] buffer = new byte[1024];
        int bytesRead;
        while ((bytesRead = inputStream.read(buffer)) != -1) {
            cipherOutputStream.write(buffer, 0, bytesRead);
        }
        cipherOutputStream.close();
        inputStream.close();
    }
    //AES解密文件
    public static void decryptFile(String inputFile, String outputFile, SecretKey key)
            throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException,
            IOException {
        Cipher cipher = Cipher.getInstance(CIPHER_TRANSFORMATION);
        cipher.init(Cipher.DECRYPT_MODE, key);
        FileInputStream inputStream = new FileInputStream(inputFile);
        CipherInputStream cipherInputStream = new CipherInputStream(inputStream, cipher);
        FileOutputStream outputStream = new FileOutputStream(outputFile);
        byte[] buffer = new byte[1024];
        int bytesRead;
        while ((bytesRead = cipherInputStream.read(buffer)) != -1) {
            outputStream.write(buffer, 0, bytesRead);
        }
        outputStream.close();
        cipherInputStream.close();
    }

}
